# 这段代码的主要功能是为图像数据创建数据加载器，同时为不同的视角（例如 UAV 和卫星图像）应用特定的数据增强和预处理操作
# 最终返回数据加载器、类别名称和数据集大小信息，以供模型训练和验证使用
# 1、数据增强与预处理配置：
# 根据传入的参数 opt 的选项，动态调整数据增强和预处理管道，包括：
# 旋转与裁剪 (RotateAndCrop)。
# 随机仿射变换 (RandomAffine)。
# 随机水平翻转 (RandomHorizontalFlip)。
# 随机擦除 (RandomErasing)。
# 色彩抖动 (ColorJitter)。
# 自动增强策略 (ImageNetPolicy)。
# 数据增强策略分为针对 UAV 图像 和 卫星图像 的不同组合，同时为验证数据集定义了单独的预处理。
# 2、数据转换（Transforms）流水线：
# 为训练（train）、验证（val）、以及卫星图像（satellite）分别定义了独立的预处理与增强流水线，并将每个流水线封装为 transforms.Compose 对象。
# 3、数据集加载：
# 使用 Dataloader_University 类自定义数据集加载逻辑。
# 结合了基于样本数量的自定义采样器 Sampler_University，支持多样本的采样逻辑。
# 4、返回值：
# dataloaders: 数据加载器，用于生成批量数据。
# class_names: 类别名称。
# dataset_sizes: 每个数据子集的大小（乘以样本数量 sample_num）

# 导入所需模块
from torchvision import transforms  # 图像增强和变换工具
from .Dataloader_University import Sampler_University, Dataloader_University, train_collate_fn  # 自定义数据加载相关类
from .autoaugment import ImageNetPolicy  # 自动增强策略
import torch
from .queryDataset import RotateAndCrop, RandomCrop, RandomErasing  # 自定义增强方法

# 定义创建数据集和数据加载器的函数
def make_dataset(opt):
    # 定义用于训练和卫星图像的增强列表
    transform_train_list = []
    transform_satellite_list = []

    # 根据 opt 参数启用 RotateAndCrop 增强，用于旋转和裁剪图像
    if "uav" in opt.rr:
        transform_train_list.append(RotateAndCrop(0.5))  # UAV 图像
    if "satellite" in opt.rr:
        transform_satellite_list.append(RotateAndCrop(0.5))  # 卫星图像

    # 添加基础增强操作，包括调整大小、填充和随机水平翻转
    transform_train_list += [
        transforms.Resize((opt.h, opt.w), interpolation=3),  # 调整到指定大小
        transforms.Pad(opt.pad, padding_mode='edge'),  # 边缘填充
        transforms.RandomHorizontalFlip(),  # 随机水平翻转
    ]
    transform_satellite_list += [
        transforms.Resize((opt.h, opt.w), interpolation=3),
        transforms.Pad(opt.pad, padding_mode='edge'),
        transforms.RandomHorizontalFlip(),
    ]

    # 定义验证集的增强（仅调整大小）
    transform_val_list = [
        transforms.Resize(size=(opt.h, opt.w), interpolation=3),  # 调整大小
    ]

    # 根据 opt 参数添加随机仿射变换
    if "uav" in opt.ra:
        transform_train_list += [transforms.RandomAffine(180)]  # UAV 图像
    if "satellite" in opt.ra:
        transform_satellite_list += [transforms.RandomAffine(180)]  # 卫星图像

    # 根据 opt 参数添加随机擦除
    if "uav" in opt.re:
        transform_train_list += [RandomErasing(probability=opt.erasing_p)]  # UAV 图像
    if "satellite" in opt.re:
        transform_satellite_list += [RandomErasing(probability=opt.erasing_p)]  # 卫星图像

    # 根据 opt 参数添加颜色抖动
    if "uav" in opt.cj:
        transform_train_list += [transforms.ColorJitter(brightness=0.5, contrast=0.1, saturation=0.1, hue=0)]  # UAV 图像
    if "satellite" in opt.cj:
        transform_satellite_list += [transforms.ColorJitter(brightness=0.5, contrast=0.1, saturation=0.1, hue=0)]  # 卫星图像

    # 如果启用了自动增强（DA），在增强列表前添加 ImageNetPolicy
    if opt.DA:
        transform_train_list = [ImageNetPolicy()] + transform_train_list

    # 最后的通用增强操作，包括转换为张量和标准化
    last_aug = [
        transforms.ToTensor(),  # 转换为张量
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # 标准化（基于 ImageNet 的均值和标准差）
    ]
    transform_train_list += last_aug
    transform_satellite_list += last_aug
    transform_val_list += last_aug

    # 打印最终的增强列表（用于调试）
    print(transform_train_list)
    print(transform_satellite_list)

    # 将训练、验证和卫星图像的增强组合到字典中
    data_transforms = {
        'train': transforms.Compose(transform_train_list),
        'val': transforms.Compose(transform_val_list),
        'satellite': transforms.Compose(transform_satellite_list)
    }

    # 使用自定义的数据加载器加载数据集
    image_datasets = Dataloader_University(
        opt.data_dir, transforms=data_transforms)  # 数据目录和增强配置
    # 定义采样器，用于控制批次中的样本分布
    samper = Sampler_University(
        image_datasets, batchsize=opt.batchsize, sample_num=opt.sample_num)
    # 创建数据加载器，支持多线程加载和自定义批处理逻辑
    dataloaders = torch.utils.data.DataLoader(image_datasets, batch_size=opt.batchsize,
                                              sampler=samper, num_workers=opt.num_worker, pin_memory=True, collate_fn=train_collate_fn)
    # 计算数据集大小（样本数量 × 采样数量）
    dataset_sizes = {x: len(image_datasets) * opt.sample_num for x in ['satellite', 'drone']}
    # 获取类别名称
    class_names = image_datasets.cls_names

    # 返回数据加载器、类别名称和数据集大小
    return dataloaders, class_names, dataset_sizes

